#!/user/bin/python
# -*- coding:utf-8 -*-

from cassandra.cluster import Cluster
from cassandra.auth import PlainTextAuthProvider
from WindPy import *
import datetime
import time

# 根据日期获取票池数组
def read_wind_stock_pool_data(str_date):
    w.start()
    res = w.wset('sectorconstituent', 'date='+str_date+';sectorid=a001010100000000;field=wind_code')
    if res.ErrorCode != 0:
        print("读取票池数据失败!")
        return
    return res.Data[0]

# 根据stock_id读取固定时间段内的收盘价
def read_wind_close(stock_id,str_begin_date, str_end_date):
    w.start()
    # resDailyClose = w.wsd(oneStockID, 'close', startDateStr, endDateStr, 'Fill=Previous;PriceAdj=' + fillType)
    res_daily_close = w.wsd(stock_id, 'close', str_begin_date, str_end_date, 'Fill=Previous;PriceAdj=' + 'F')
    times = res_daily_close.Times
    data = res_daily_close.Data[0]
    str_times = []
    for t in times:
        str_times.append(t.strftime('%Y-%m-%d'))
    result = {'str_times': str_times, 'times': times, 'data': data}
    print(result)
    return result


def import_daily_close_data(stock_id_array, str_begin_date, str_end_date):
    #stock_id = '000001.SZ'
    conn = get_connection()
    conn.set_keyspace('stock')
    prepared_statement = conn.prepare("INSERT INTO  daily_close(stock,datetime,stock_value) VALUES(?, ?, ?)")
    for stock_id in stock_id_array:
        result = read_wind_close(stock_id, '2012-01-01', '2012-01-31')
        str_times = result.get('str_times')
        times = result.get('times')
        datas = result.get('data')
        data_len = len(datas)
        # insert into daily_close(stock,datetime,stock_value) values('000001.SZ','2017-04-05',964.82351);
        for i in range(0, data_len):
            print((stock_id, times[i], datas[i]))
            conn.execute_async(prepared_statement, (stock_id, times[i], datas[i]))



# 根据开始日期和结束日期生成日期数组
# x = gen_data_range(datetime.date(2012, 1, 1), datetime.date(2017,1,1))
# print(x)
def gen_data_range(begin_date, end_date):
    d = begin_date
    x = []
    delta = datetime.timedelta(days=1)
    while d <= end_date:
        x.append(d.strftime('%Y-%m-%d'))
        d += delta
    return x

# 获取cassandra数据库连接
def get_connection():
    # 口令
    auth_provider = PlainTextAuthProvider(username='root', password='123456')
    # 创建集群
    cluster = Cluster(contact_points=['192.168.100.104'], auth_provider=auth_provider)
    # 获取数据库连接
    return cluster.connect()


# 初始化数据库表
def init_cassandra_db(conn):
    # 创建数据库stock
    conn.execute("CREATE KEYSPACE IF NOT EXISTS stock WITH REPLICATION={'class':'SimpleStrategy','replication_factor':1}")

    # 切换数据库
    conn.set_keyspace('stock')

    # 创建票池stock_pool表
    conn.execute("CREATE TABLE IF NOT EXISTS stock_pool(datetime timestamp primary key, stock_list list<text>, stock_size int)")

    # 创建daily_close表
    conn.execute("CREATE TABLE IF NOT EXISTS daily_close(stock text, datetime date, stock_value double, primary key(stock,datetime))")
    print("cassandra init successfully!!!!!!")

# 导入票池数据
def import_stock_pool_data(conn):
    date_range = gen_data_range(datetime.date(2012, 1, 1), datetime.date(2017, 5, 8))
    conn.set_keyspace('stock')
    prepared_statement = conn.prepare("INSERT INTO  stock_pool(datetime, stock_list, stock_size) VALUES(?, ?, ?)")
    for current_date in date_range:
        data = read_wind_stock_pool_data(current_date)
        t = time.strptime(current_date, '%Y-%m-%d')
        param_current_date = datetime.datetime(*t[:3])
        print(param_current_date)
        conn.execute_async(prepared_statement, (param_current_date, data, len(data)))


def process_stock_pool():
    conn = get_connection()
    init_cassandra_db(conn)
    import_stock_pool_data(conn)
    conn.shutdown()

# 读取票池数据
def read_stock_pool(str_begin_date, str_end_date):
    conn = get_connection()
    conn.set_keyspace("stock")
    sql = "SELECT * FROM stock_pool WHERE datetime >=  '"+str_begin_date+" 00:00:00' and datetime <= '"+str_end_date+" 23:59:59' allow filtering";
    rows = conn.execute(sql)
    map = {}
    for row in rows:
        map[row.datetime.strftime('%Y-%m-%d')] = row.stock_list
        print(row.datetime.strftime('%Y-%m-%d'), row.stock_list)
    # 根据日期获取当日的票池
    # print(map.get('2012-01-01'))
    return map

# 程序运行的入口
if __name__ == '__main__':

    #str_begin_date = '2012-01-01'
    #str_end_date = '2012-01-31'
    #stock_pool_data = read_stock_pool(str_begin_date, str_end_date)
    #stock_id_array = stock_pool_data.get('2012-01-09')
    #import_daily_close_data(stock_id_array, str_begin_date, str_end_date)



    print(read_wind_stock_pool_data('201601'))



